import pandas as pd
import torch
from PIL import Image
import io
from IPython.display import display
import numpy as np
import random


def random_model(bytes_input):
    # 把 bytes 转换为 numpy 数组
    img = Image.open(io.BytesIO(bytes_input))
    print(img.size)
    img = img.resize((224, 224))
    print(img.size)
    img = img.convert('RGB')
    print(img.size)
    img_array = np.array(img)
    print('numpy array shape:', img_array.shape)
    model_input = torch.from_numpy(img_array).float()
    print('tensor ' , model_input.size())
    model_input = model_input.unsqueeze(0)
    print('tensor ', model_input.size())
    return random.randint(0, 9)


# 读取 Parquet 文件
df = pd.read_parquet('train.parquet', engine='pyarrow')
# 显示第一张图片bytes数据及标签
img = df.iloc[0]['image']['bytes']
label = df.iloc[0]['label']
image = Image.open(io.BytesIO(img))
print(image.size)
print('image:')
display(image)
print('label:', label)
print('random model prediction:', random_model(img))

